#include "sweSolver.hpp"
#include "sweFunctions.hpp"
#include "macros.hpp"

namespace hms {

/** @brief Calculates the time step according to the maximum stable time step
 * (CFL) and the maximum Courant number */
template<typename MeshType>
auto SWESolver<MeshType>::calculateTimeStepImp() const -> scalar {
	scalar step { std::numeric_limits<scalar>::max() };

	/* the "courant factor" is 4 * cell area over cell circumference. 
	 * For different mesh types, the optimal representation differs and is
	 * given by this lambda */
	auto getCourantFactor = [&](
		[[maybe_unused]] Index iBeg,
		[[maybe_unused]] Index colBeg,
		[[maybe_unused]] Index nCols,
		[[maybe_unused]] Index j
	){
		if constexpr ( std::is_same_v<MeshType, UniMesh> ){

			auto cell { this->mesh().cell(0,0) };
			return 4 * cell.area() / cell.circumf();

		} else if constexpr ( std::is_same_v<MeshType, RectMesh> ){

			Eigen::VectorBlock<const ArrayXs> dx {
				this->mesh().dx().segment(iBeg, nCols)
			};
			scalar dy { this->mesh().dy()[j] };
			return 4 * (dx * dy / ( 2 * (dx + dy) )).transpose();

		} else if constexpr (
			std::is_same_v<MeshType, StructMesh> ||
			std::is_same_v<MeshType, Mesh>
		){
			const auto& cells { this->mesh().cells() };
			Eigen::VectorBlock<const ArrayXs>
				areas    { cells.areas   ().segment(colBeg, nCols) },
				circumfs { cells.circumfs().segment(colBeg, nCols) };

			return 4 * (areas / circumfs).transpose();
		}
	};
	
	auto blockwiseMinStep = [&](
		Index iBeg, Index colBeg, Index nCols, Index j
	) -> scalar {
		/* using Eigen::Map here was not faster than middleCols */
		return ( getCourantFactor(iBeg, colBeg, nCols, j) / (
			v().values().middleCols(colBeg, nCols).colwise().norm() +
			(d().values().middleCols(colBeg, nCols) * phys::g).sqrt()
		) ).minCoeff();
	};

	if constexpr ( std::is_base_of_v<StructMeshBase<MeshType>, MeshType> ){
		static constexpr Index ng {1};
		Index
			nx { this->mesh().nCellsX() },
			ny { this->mesh().nCellsY() },
			cellsPerRow { nx + 2*ng };

/** @note First note:
 * Two alternatives: Just going through the mesh row by row vs doing top
 * and bottom boundaries first and then having longer segments afterwards,
 * which is the commented out section below the loop.
 * Tested on 200x200, 500x500 and 1000x1000.
 * On 200x200 it's 1.2% vs 1.6%, on 500x500 it's 1.1% vs 1.15% and on 1000x1000
 * it's 1.5% vs 1.45%.
 * Therefore, on large local meshes, it is slightly slower to go row by row,
 * but at the cost of higher code complexity and much slower execution on
 * smaller meshes.
 * Recommendation: Simpler version, row by row.
 * 
 * Second note:
 * added another simpler version which only takes domain cells into account.
 * */
		// if constexpr ( std::is_same_v<MeshType, RectMesh> ){
			/* take all cells except corner cells into account */
			// OMP_PRAGMA( parallel for reduction (min:step) )
			// for ( int j=0; j<ny+2*ng; ++j ){
			// 	/* loop over all rows but skip corners */
			// 	Index
			// 		iBeg   = (j>=ng && j<ng+ny) ? 0 : ng,
			// 		colBeg = j*cellsPerRow + iBeg,
			// 		nCols  =   cellsPerRow - 2*iBeg;
			// 	step = std::min( step, blockwiseMinStep( iBeg, colBeg, nCols, j ) );
			// }

			/* Take only domain cells into account */
			OMP_PRAGMA( parallel for reduction (min:step) )
			for ( int j = ng; j<ny+ng; ++j){
				Index colBeg { j*cellsPerRow + ng };
				step = std::min( step, blockwiseMinStep( ng, colBeg, nx, j ) );
			}
		// } else {
		// 	/* bottom and top ghosts in these sections, because the rest
		// 	 * is just one long contiguous segment of cells */
		// 	OMP_PRAGMA( parallel sections reduction (min:step) )
		// 	{
		// 		OMP_PRAGMA( section )
		// 		for ( int j=0; j<ng; ++j )
		// 			step = std::min( step,
		// 				blockwiseMinStep(ng, j*cellsPerRow + ng, nx, j)
		// 			);
		// 		OMP_PRAGMA( section )
		// 		for ( int j=ng+ny; j<ny+2*ng; ++j )
		// 			step = std::min( step,
		// 				blockwiseMinStep(ng, j*cellsPerRow + ng, nx, j)
		// 			);
		// 	}
		// 	OMP_PRAGMA( parallel reduction (min:step) )
		// 	{
		// 		auto [colBeg, nCols] = ompSegment( cellsPerRow*ny );
		// 		/* the segments are zero-based, but the bottom ghost cells
		// 		 * are already covered. */
		// 		colBeg += ng*cellsPerRow;

		// 		step = std::min( step, blockwiseMinStep(0, colBeg, nCols, 0) );
		// 	}
		// }
	} else {
		OMP_PRAGMA( parallel reduction (min:step) )
		{

		auto [colBeg, nCols] = ompSegment( this->mesh().nCells() );
		
		step = std::min(
			step, blockwiseMinStep(Index{}, colBeg, nCols, Index{})
		);

		}
	}

	// /* non-blockwise version, slower */
	// if constexpr ( std::is_base_of_v<StructMeshBase<MeshType>, MeshType> ){
	// 	Index
	// 		ng { this->mesh().nGhostLayers() },
	// 		nx { this->mesh().nCellsX() },
	// 		ny { this->mesh().nCellsY() };
	// 	OMP_PRAGMA( parallel for collapse(2) reduction (min:step) )
	// 	for ( Index j = -ng; j < ny+ng; ++j ){
	// 		for ( Index i = -ng; i < nx+ng; ++i ){
	// 			/* loop over all cells but skip corners */
	// 			if ( ( j<0 || j>=ny ) && ( i<0 || i>=nx ) )
	// 				continue;
	// 			step = std::min( step,
	// 				courantTimeStep( this->mesh().cell(i, j), v(), d() )
	// 			);
	// 		}
	// 	}
	// } else {
	// 	OMP_PRAGMA( parallel for reduction (min:step) )
	// 	for ( Index i = 0; i<this->mesh().nCells(); ++i )
	// 		step = std::min( step,
	// 			courantTimeStep( this->mesh().cell(i), v(), d() )
	// 		);
	// }
	
	return this->settings().maxCourant * step;
}

template class SWESolver<Mesh>;
template class SWESolver<UniMesh>;
template class SWESolver<RectMesh>;
template class SWESolver<StructMesh>;

} // end namespace hms
